"""
Broker encapsulates SQLAlchemy functionality in a simple intuitive class
Written By: Daniel Ortiz
Date: November 3 2007
Version: 0.0
"""

from sqlalchemy import Table, Integer, String, Column, ForeignKey, MetaData, Boolean, DateTime, create_engine, Binary, Text
from sqlalchemy.orm import mapper, relation, sessionmaker, clear_mappers
from mapclass import User, Owner, Mailbox, Mail, Domain, Address, AddressBook, Control, Attachment, Folder
from sasync.database import AccessBroker, transact
import datetime
from twisted.internet import defer

class EmptyObjectException(Exception):
    """Exception occurs when object does not exist in the database"""
    pass

class DatabaseCollision(Exception):
    """Exception called when object found in the database attemps to be reinserted"""
    pass

DB_URL = "sqlite:///database.db"

class ABroker(AccessBroker):
    """Access Broker for the database"""

    def __getattribute__(self, name):
        # For debugging, to see what session is stored
        result = AccessBroker.__getattribute__(self, name)
        return result
    
    def startup(self):
        def gotTables(null):
            mapper(AddressBook, self.addressbook, properties = {
                'address' : relation(Address,
                                     backref = 'addressbook',
                                     cascade = "all, delete, delete-orphan"),
                }
            )
            mapper(Address, self.address)
            mapper(Mailbox, self.mailbox, properties = {
                'mail' : relation(Mail,
                                  backref = 'mailbox',
                                  cascade = "all, delete, delete-orphan"),
                }
            )
            mapper(Mail, self.mail)
            
        d = self.table('addressbook',
            Column('id', Integer, primary_key = True),
            Column("alias", String(255)))
        d.addCallback(lambda _: self.table('address',
            Column('id', Integer, primary_key = True),
            Column('abook_id', Integer, ForeignKey('addressbook.id')),
            Column('fullname', String(255)),
            Column('email', String(255))))
        d.addCallback(lambda _: self.table('mailbox', 
            Column('id', Integer, primary_key = True),
            Column('alias', String(255)),
            Column('size', Integer),
            Column('username', String(255)),
            Column('domain', String(255)),
            Column('password', String(255)),
            Column('dtime', Integer),
            Column('protocol', String(255)),
            Column('iserver', String(255)),
            Column('iport', Integer),
            Column('issl', Boolean),
            Column('oserver', String(255)),
            Column('oport', Integer),
            Column('ossl', Boolean)))
        d.addCallback(lambda _: self.table('mail',
            Column('id', Integer, primary_key = True),
            Column('orig_date', Text),
            Column('from_who', Text),
            Column('sender', Text),
            Column('uid', Integer),
            Column('flags', Text),
            Column('local_flags', Text),
            Column('reply_to', Text),
            Column('body', Text),
            Column('to', Text),
            Column('cc', Text),
            Column('bcc', Text),
            Column('subject', Text),
            Column('encoding', Text),
            Column('content_type', Text),
            Column('charset', Text),
            Column('folder', Text),
            Column('mailbox_id', Integer, ForeignKey('mailbox.id'))))
     
        d.addCallback(gotTables)
        return d
    
    @transact
    def DeleteMail(self, mailbox, mail):
        m = self.GetMail(mailbox, mail)
        self.session.delete(m)
    
    @transact
    def GetMail(self, mailbox, mail):
        return self.session.query(Mail).filter_by(**mail.kwargs()).filter(Mail.mailbox.has(**mailbox.kwargs()))[0]
    
    @transact
    def SaveMailList(self, mailList, mailbox):
        dobj = []
        box = self.GetMailbox(mailbox)
        for m in mailList:
            mail = Mail(*m.get())
            box.mail.append(mail)
            self.session.save(mail)
            dobj.append(mail) 
            print mail.body, mailbox           
        return dobj
    
    @transact
    def DeleteMailList(self, maillist, mailbox):
        for m in maillist:
            self.DeleteMail(mailbox, m)
            
    @transact
    def UpdateMail(self, old_mail, new_mail, mailbox):
        m = self.GetMail(mailbox, old_mail)
        m.update(new_mail)
    
    @transact
    def GetMailbox(self, mailbox):
        return self.session.query(Mailbox).filter_by(**mailbox.kwargs())[0]
    
    @transact
    def UpdateMailbox(self, old_mailbox, new_mailbox):
        box = self.GetMailbox(old_mailbox)
        box.update(new_mailbox)
        print box.alias
        
    @transact
    def SaveMailbox(self, mailbox):
        self.session.save(Mailbox(*mailbox.get()))
        
    @transact
    def SaveAddressBookList(self, ABList):
        for AB in ABList:
            if isinstance(AB, basestring):
                self.session.save(AddressBook(AB))
            else:
                self.session.save(AddressBook(*AB))

    @transact
    def SaveAddressBook(self, addressbook):
        self.session.save(AddressBook(addressbook))

    @transact
    def UpdateAddressBook(self, oldbook, newbook):
        book = self.GetAddressBook(oldbook)[0]
        book.alias = newbook
        
    @transact
    def DeleteAddressBook(self, addressbook):
        book = self.GetAddressBook(addressbook)
        self.session.delete(book[0])

    @transact
    def GetAddressBook(self, addressbook):
        return self.session.query(AddressBook).filter_by(alias = addressbook)

    @transact
    def GetAddressBookNames(self):
        return [m.alias for m in self.session.query(AddressBook)]
    
    @transact
    def SaveAddressList(self, addresslist, addressbook):
        book = self.GetAddressBook(addressbook)
        for atup in addresslist:
            address = Address(*atup)
            book[0].address.append(address)
            self.session.save(address)

    @transact
    def SaveAddress(self, atup, addressbook):
        book = self.GetAddressBook(addressbook)
        address = Address(*atup)
        book[0].address.append(address)
        self.session.save(address)
        
    @transact
    def UpdateAddress(self, book, oaddress, naddress):
        a = self.session.query(Address).filter_by(email = oaddress[0], fullname = oaddress[1]).\
            filter(Address.addressbook.has(AddressBook.alias == book))[0]
        a.email = naddress[0]
        a.fullname = naddress[1]

    @transact
    def DeleteAddress(self, book, address):
        a = self.session.query(Address).filter_by(email = address[0], fullname = address[1]).\
            filter(Address.addressbook.has(AddressBook.alias == book))[0]
        self.session.delete(a)
        
    @transact
    def GetABAddress(self, addressbook):
        return self.session.query(Address).filter(Address.addressbook.has(AddressBook.alias == addressbook))

    @transact
    def GetABAddressAttr(self, addressbook, attr):
        address = self.GetABAddress(addressbook)
        return [getattr(m, attr) for m in address]
        
    @transact
    def GetData(self):
        ret, ret2 = {},{}
        mailbox = self.session.query(Mailbox)
        addressbook = self.session.query(AddressBook)
        for box in mailbox:
            ret[box] = []
            for mail in box.mail:
                ret[box].append(mail)
        for book in addressbook:
            ret2[book.alias] = []
            for address in book.address:
                ret2[book.alias].append((address.email, address.fullname))
        return ret, ret2
    
    @transact
    def Close(self):
        self.session.flush()
        self.session.close()
        clear_mappers()
        self.shutdown()
#        
#class AsynchronousMailBroker(AccessBroker):
#    """
#    Asynchronous mail broker, which works with deferred"
#    """
#    def startup(self):
#        addressbook = self.table('addressbook',
#            Column('id', Integer, primary_key = True),
#            Column("alias", String(255)))
#        address = self.table('address',
#            Column('id', Integer, primary_key = True),
#            Column('abook_id', Integer, ForeignKey('addressbook.id')),
#            Column('fullname', String(255)),
#            Column('email', String(255)))
#        mailbox = self.table('mailbox', 
#            Column('id', Integer, primary_key = True),
#            Column('alias', String(255)),
#            Column('size', Integer),
#            Column('username', String(255)),
#            Column('domain', String(255)),
#            Column('password', String(255)),
#            Column('dtime', Integer),
#            Column('protocol', String(255)),
#            Column('iserver', String(255)),
#            Column('iport', Integer),
#            Column('issl', Boolean),
#            Column('oserver', String(255)),
#            Column('oport', Integer),
#            Column('ossl', Boolean))
#        mail = self.table('mail',
#            Column('id', Integer, primary_key = True),
#            Column('orig_date', Text),
#            Column('from_who', Text),
#            Column('sender', Text),
#            Column('uid', Integer),
#            Column('flags', Text),
#            Column('local_flags', Text),
#            Column('reply_to', Text),
#            Column('body', Text),
#            Column('to', Text),
#            Column('cc', Text),
#            Column('bcc', Text),
#            Column('subject', Text),
#            Column('encoding', Text),
#            Column('content_type', Text),
#            Column('charset', Text),
#            Column('folder', Text),
#            Column('mailbox_id', Integer, ForeignKey('mailbox.id')))
#        
#        return defer.DeferredList((mailbox, mail, address, addressbook))
#    
#    @transact
#    def DeleteMail(self, mail):
#        self.mail.delete().execute(mail.kwargs())
#        
#    @transact
#    def DeleteAddress(self, address):
#        self.address.delete().execute(address.kwargs())
#        
#    @transact
#    def DeleteAddressBook(self, addressbook):
#        self.addressbook.delete().excute(addressbook.kwargs())
#    
#    @transact
#    def GetMail(self, mail):
#        return Mail(*self.mail.select().execute(mail.kwargs()).fetchone())
#    @transact
#    def SaveMailList(self, mail_list):
#        self.mail.insert().execute([mail.kwargs() for mail in mail_list])
#    
#    @transact
#    def SaveMail(self, mail):
#        self.mail.insert().execute(mail.kwargs())
#    
#    @transact
#    def UpdateMail(self, old_mail, new_mail, mailbox):
#        def Update(conn):
#            conn.execute(self.mail.update(self.mail.c.id == old_mail.id), new_mail.kwargs())
#        self.connect().addCallback(Update)
#        
#    @transact 
#    def UpdateMailbox(self, old_mailbox, new_mailbox):
#        def Update(conn):
#            conn.execute(self.mailbox.update(self.mailbox.c.id == old_mailbox.id), new_mailbox.kwargs())
#        self.connect().addCallback(Update)
#        
#    @transact
#    def UpdateAddress(self, old_address, new_address):
#        def Update(conn):
#            conn.execute(self.address.update(self.address.c.id == old_address.id), new_address.kwargs())
#        self.connect().addCallback(Update)
#        
#    @transact
#    def ModifyAddressBook(self, old_book, new_book):
#        def Update(conn):
#            conn.execute(self.addressbook.update(self.addressbook.c.id == old_book.id), new_book.kwargs()) 
#        self.connect().addCallback(Update)
#        
#    @transact
#    def GetMailbox(self, mailbox):
#        def Update(conn):
#            c = conn.execute(self.mailbox.select(), mailbox.kwargs()).fetchone()
#            return Mailbox(*c)
#        m = self.connect().addCallback(Update)
#        return m.result
#        #return Mailbox(*self.mailbox.select().execute(mailbox.kwargs()).fetchone())
#    
#    @transact
#    def GetAddressBook(self, book):
#        def Update(conn):
#            for m in conn.execute(self.addressbook.select(), book.kwargs()):
#                print "First",  AddressBook(*m)
#            c = conn.execute(self.addressbook.select(), book.kwargs()).fetchone()
#            return AddressBook(*c)
#        m = self.connect().addCallback(Update)
#        print book, m.result
#        return m.result
#    
#    @transact
#    def GetIDAddressBook(self, id):
#        mbook = self.addressbook.select().execute().fetchall()
#        for m in mbook:
#            if m[0] == id:
#                return AddressBook(*m)
#    
#    @transact
#    def SaveMailbox(self, mailbox):
#        def Update(conn):
#            i = self.mailbox.insert()
#            conn.execute(i, [mailbox.kwargs()])
#        self.connect().addCallback(Update)
#        
#    @transact
#    def SaveAddressBook(self, book):
##        def Update(conn):
##            conn.execute(self.addressbook.insert(), [book.kwargs()])
##            #print self.addressbook.select().execute().fetchall()
##        self.connect().addCallback(Update)
#        return self.addressbook.insert().execute(alias = book.alias).fetchall()
#        
#    @transact
#    def DeleteAll(self):
#        self.mail.delete().execute()
#        self.mailbox.delete().execute()
#        
#    @transact
#    def GetData(self):
#        ret, ret2 = {}, {}
#        mailbox = self.mailbox.select().execute().fetchall()
#        mail = self.mail.select().execute().fetchall()    
#        address = self.address.select().execute().fetchall()
#        addressbook = self.addressbook.select().execute().fetchall()
#        for box in mailbox:
#            id = box[0]
#            print id
#            ret[id] = {'MAILBOX': Mailbox(*box), 'MAIL' : []}
#        for m in mail:
#            ret[m.mailbox_id]['MAIL'].append(Mail(*m))
#        for book in addressbook:
#            print book[0]
#            ret2[AddressBook(*book)] = []
#        for a in address:
#            tbook = None
#            for book in ret2.keys():
#                if book.id == a.abook_id:
#                    tbook = book
#                    break
#            ret2[tbook].append(Address(*a))
#        return ret, ret2

class MailBroker(object):
    """
    New class acts as an interface to an SQLAlchemy database. Functions are
    defined as ActionClass. In other words, InsertUser is a function which
    inserts(Action) a User(Class) into the database. The definitions of class-
    methods fits this paradigm.
    """
    def __init__(self, dbname = 'mail.db', path='.'):
        """
        Constructor for the broker class initializes all SQLAlchemy tables
        and their mapper representations.

        dbname - name of the database
        path - location of the database
        """

        # Setup metadata and engine
        engine = create_engine('sqlite://./%s' % dbname)
        metadata = MetaData()

        # Tables in the mail application
        owner_table = Table('owner', metadata,
            Column('id', Integer, primary_key = True),
            Column('name', String(255)),
            Column('password', String(255)))
        user_table = Table('users', metadata,
            Column('id', Integer, primary_key = True),
            Column('username', String(255)),
            Column('password', String(255)),
            Column('owner_id', Integer, ForeignKey('owner.id')))
        control_table = Table('control', metadata,
            Column('id', Integer, primary_key = True),
            Column('multiple_users', Boolean),
            Column('sc_arch', Boolean),
            Column('password', Boolean),
            Column('lmailbox', String(255)),
            Column('owner_id', Integer, ForeignKey('owner.id')))
        addressbook_table = Table('addressbook', metadata,
            Column('id', Integer, primary_key = True),
            Column('private', Boolean),
            Column('alias', String(255)),
            Column('user_id', Integer, ForeignKey('users.id')))
        address_table = Table('address', metadata,
            Column('id', Integer, primary_key = True),
            Column('email', String),
            Column('fullname', String),
            Column('addressbook_id', Integer, ForeignKey('addressbook.id')))
        mailbox_table = Table('mailbox', metadata,
            Column('id', Integer, primary_key = True),
            Column('alias', String(255)),
            Column('private', Boolean),
            Column('size', Integer),
            Column('users_id', Integer, ForeignKey('users.id')))
        domain_table = Table('domain', metadata,
            Column('id', Integer, primary_key = True),
            Column('username', String(255)),
            Column('domain', String(255)),
            Column('password', String(255)),
            Column('dtime', Integer),
            Column('protocol', String(255)),
            Column('iserver', String(255)),
            Column('iport', Integer),
            Column('issl', Boolean),
            Column('oserver', String(255)),
            Column('oport', Integer),
            Column('ossl', Boolean),
            Column('mail', String(255)),
            Column('mailbox_id', Integer, ForeignKey('mailbox.id')))
        mail_table = Table('mail', metadata,
            Column('id', Integer, primary_key = True),
            Column('orig_date', DateTime),
            Column('from_who', String),
            Column('body', String),
            Column('state', String),
            Column('sender', String),
            Column('uid', Integer),
            Column('flags', String),
            Column('reply_to', String),
            Column('to', String),
            Column('cc', String),
            Column('bcc', String),
            Column('message_id', String),
            Column('in_reply_to', String),
            Column('references', String),
            Column('subject', String),
            Column('encoding', String),
            Column('mailbox_id', Integer, ForeignKey('mailbox.id')))
        attachment_table = Table('attachment', metadata,
            Column('id', Integer, primary_key = True),
            Column('title', String),
            Column('data', String),
            Column('size', Integer),
            Column('encoding', String),
            Column('alternative', Boolean),
            Column('mixed', Boolean),
            Column('mail_id', Integer, ForeignKey('mail.id')))
        folder_table = Table('folder', metadata,
            Column('id', Integer, primary_key = True),
            Column('alias', String),
            Column('mailbox_id', Integer, ForeignKey('mailbox.id')))

        # Create all of the tables in the database
        metadata.create_all(engine)

        # Map all of the data.
        mapper(Owner, owner_table, properties = {
            'users':relation(User, backref = 'owner'),
            'control': relation(Control, backref = 'owner'),
            }
        )
        mapper(User, user_table, properties = {
            'mailbox':relation(Mailbox,
                               backref = 'users',
                               cascade = 'all, delete, delete-orphan'),
            'addressbook': relation(AddressBook,
                                    backref = 'users',
                                    cascade = 'all, delete, delete-orphan'),
            }
        )
        mapper(Mailbox, mailbox_table, properties = {
            'domain': relation(Domain,
                               backref = 'mailbox',
                               cascade = 'all, delete, delete-orphan'),
            'mail': relation(Mail,
                             backref = 'mailbox',
                             cascade = 'all, delete, delete-orphan'),
            'folder': relation(Folder,
                             backref = 'mailbox',
                             cascade = 'all, delete, delete-orphan'),
            }
        )
        mapper(AddressBook, addressbook_table, properties = {
            'address': relation(Address,
                                backref = 'addressbook',
                                cascade = 'all, delete, delete-orphan')
            }
        )
        mapper(Mail, mail_table, properties = {
            'attachment': relation(Attachment,
                                   backref = 'mail',
                                   cascade = 'all, delete, delete-orphan')
            }
        )
        mapper(Attachment, attachment_table)
        mapper(Control, control_table)
        mapper(Domain, domain_table)
        mapper(Address, address_table)
        mapper(Folder, folder_table)

        # Create the session
        Session = sessionmaker(bind = engine, autoflush = True, transactional = True)
        self.session = Session()

    # VALIDATE functions: Test whether an object is valid
    # def ValidateClass(self, SampleClass)
    # Class - the class you would like to validate
    # SampleClass - an object of that class you would like to validate
    def ValidateUser(self, iUser):
        if isinstance(iUser, User):
            user = self.session.query(User).filter_by(username = iUser.username)
            if user:
                try:
                    return user[0].password == iUser.password
                except IndexError:
                    return False
            else:
                return False
        else:
            return False

    # GET Functions: Get an object from the database
    # def GetClass(self, SampleClass)
    # Class - the class you would like to get
    # SampleClass - an object of the above Class
    # Return: sqlalchemy.Query
    # def Get(ParentClass)Class(self, SampleClass, ParentClass...)
    # Class - the class you would like to get
    # SampleClass - and object of the class
    # ParentClass - a parent object, which will filter the results
    def GetOwner(self):
        return self.session.query(Owner)

    def GetUser(self):
        return self.session.query(User)

    def GetSuperUser(self):
        """Function gets the user that is the same instance as the owner"""
        owner = self.GetOwner()
        if owner:
            return self.session.query(User).filter_by(username = owner[0].name,
                                                      password = owner[0].password)
        else:
            raise EmptyObjectException('Owner does not exists, therefore there is no superuser')

    def GetControl(self):
        return self.session.query(Control)

    def GetMailbox(self):
        return self.session.query(Mailbox)

    def GetUserMailbox(self, iUser):
        return self.session.query(Mailbox).filter(Mailbox.users.has(User.username == iUser.username))

    def GetAddressBook(self):
        return self.session.query(AddressBook)

    def GetUserAddressBook(self, iUser):
        return self.session.query(AddressBook).\
               filter(AddressBook.users.has(User.username == iUser.username))

    def GetMail(self):
        return self.session.query(Mail)

    def GetMailboxMail(self, iMailbox):
        return self.session.query(Mail).filter(Mail.mailbox.has(Mailbox.alias == iMailbox.alias))

    def GetDomain(self):
        return self.session.query(Domain)

    def GetMailboxDomain(self, iMailbox):
        return self.session.query(Domain).filter(Domain.mailbox.has(Mailbox.alias == iMailbox.alias))

    def GetMailboxFolder(self, iMailbox):
        return self.session.query(Folder).filter(Folder.mailbox.has(Mailbox.alias == iMailbox.alias))
    
    def GetAddress(self):
        return self.session.query(Address)

    def GetAddressBookAddress(self, iAddressBook):
        return self.session.query(Address).filter(Address.addressbook.has(AddressBook.alias == iAddressBook.alias))

    # DELETE FUNCTIONS: Deletes an object from the database.
    # def DeleteClass(self, SampleClass)
    # Class - the class type you would like to delete
    # SampleClass - an object of the above class to compare to the database
    # Return - sqlalchemy.Query
    def DeleteOwner(self, iOwner):
        return self.session.delete(iOwner)

    def DeleteUser(self, iUser):
        return self.session.delete(iUser)

    def DeleteControl(self, iControl):
        return self.session.delete(iControl)

    def DeleteDomain(self, iDomain):
        return self.session.delete(iDomain)

    def DeleteMail(self, iMail, iMailbox):
        """Load a piece of mail into the session that appears like iMail, and delete it"""
        m = self.GetMailboxMail(iMailbox)
        for val in m:
            if val == iMail:
                return self.session.delete(val)

    def DeleteMailbox(self, iMailbox):
        return self.session.delete(iMailbox)

    def DeleteAddressBook(self, iAddressBook):
        return self.session.delete(iAddressBook)

    def DeleteAddress(self, iAddress):
        return self.session.delete(iAddress)

    def DeleteFolder(self, iFolder):
        return self.session.delete(iFolder)

    # HAS FUNCTIONS: Tests if the database has an object
    # def HasClass(self, SampleClass)
    # Class - the class type you would like to compare to query.
    # SampleClass - an object of the above Class to compare to the database
    # Return - Returns True if the object is in the database
    def HasOwner(self, iOwner):
        return self.session.query(Owner).filter_by(name = iOwner.name).count() != 0

    def HasUser(self, iUser):
        return self.session.query(User).filter_by(username = iUser.username).count() != 0

    def HasControl(self, iControl):
        return self.session.query(Control).count() != 0

    def HasMailbox(self, iMailbox, iUser):
        return self.session.query(Mailbox).filter_by(alias = iMailbox.alias).\
                  filter(Mailbox.users.has(User.username == iUser.username)).count() != 0

    def HasDomain(self, iDomain, iMailbox, iUser):
        return self.session.query(Domain).\
               filter(Domain.mailbox.has(Mailbox.alias == iMailbox.alias)).\
               filter(Mailbox.users.has(User.username == iUser.username)).count() != 0

    def HasMail(self, iMail, iMailbox, iUser):
        return self.session.query(Mail).\
               filter(Mail.mailbox.has(Mailbox.alias == iMailbox.alias)).\
               filter(Mailbox.users.has(User.username == iUser.username)).count() != 0

    def HasAddressBook(self, iAddressBook, iUser):
        return self.session.query(AddressBook).filter(AddressBook.alias == iAddressBook.alias).\
               filter(AddressBook.users.has(User.username == iUser.username)).count() != 0

    def HasAddress(self, iAddress, iAddressBook, iUser):
        return self.session.query(Address).filter(Address.email == iAddress.email).\
            filter(Address.addressbook.has(AddressBook.alias == iAddressBook.alias)).\
            filter(AddressBook.users.has(User.username == iUser.username)) != 0
                

    # SAVE FUNCTIONS : Insert an object into the database
    # def InsertClass(self, ParentClass):
    # Class - the object you would like to insert (User)
    # ParentClass - the parent class of this object for orm backrefences (if User,
    # then Owner would be a parent class. Only true if the object has a parent class
    def SaveOwner(self, iOwner):
        if self.HasOwner(iOwner):
            raise DatabaseCollision('Owner %s already exists in the database' % iOwner.name)
        self.session.save(iOwner)

    def SaveUser(self, iUser, iOwner):
        owner = self.session.query(Owner).filter_by(name = iOwner.name)
        if owner:
            self.session.save(iUser)
            owner[0].users.append(iUser)
        else:
            raise EmptyObjectException("Owner named %s does not exist in the database" % iOwner.name)

    def SaveControl(self, iControl, iOwner):
        owner = self.session.query(Owner).filter_by(name = iOwner.name)
        if owner:
            if not len(owner[0].control):
                self.session.save(iControl)
                owner[0].control.append(iControl)
            else:
                # It is necessary to copy over attributes, since their can
                # only be on control
                c = owner[0].control[0]
                c.multiple_users = iControl.multiple_users
                c.password = iControl.password
                c.sc_arch = iControl.sc_arch
                c.lmailbox = iControl.lmailbox
                self.session.save(c)
                self.session.save(owner[0])
        else:
            raise EmptyObjectException("Owner named %s does not exist in the database" % iOwner.name)

    def SaveMailbox(self, iMailbox, iUser):
        user = self.session.query(User).filter_by(username = iUser.username)
        if user:
            count = 0

            # If the mailbox exists in the database, just update.
            for v in user[0].mailbox:
                if v.alias == iMailbox.alias:
                    m = user[0].mailbox[count]
                    m.alias = iMailbox.alias
                    m.private = iMailbox.private
                    m.size = iMailbox.size
                    self.session.save(m)
                    self.session.save(user[0])
                    return
                count = count + 1
            user[0].mailbox.append(iMailbox)
            self.session.save(iMailbox)
        else:
            raise EmptyObjectException('User named %s does not exist in the database' % iUser.username)

    def SaveDomain(self, iDomain, iMailbox, iUser):
        mailbox = self.session.query(Mailbox).filter_by(alias = iMailbox.alias).\
                  filter(Mailbox.users.has(User.username == iUser.username))
        if mailbox:
            mailbox[0].domain.append(iDomain)
            self.session.save(iDomain)
        else:
            raise EmptyObjectException('User named %s does not have a mailbox named %s' % (iUser.username,
                                                                                            iMailbox.alias))

    def SaveFolder(self, iFolder, iMailbox, iUser):
        mailbox = self.session.query(Mailbox).filter_by(alias = iMailbox.alias).\
                  filter(Mailbox.users.has(User.username == iUser.username))
        if mailbox:
            # Only update if it is a new folder
            newFolder = True
            for v in mailbox[0].folder:
                if v.alias == iFolder.alias:
                    newFolder = False
            if newFolder:
                mailbox[0].folder.append(iFolder)
                self.session.save(iFolder)
        else:
            raise EmptyObjectException('User named %s does not have a mailbox named %s' % (iUser.username,
                                                                                            iMailbox.alias))
   
    def SaveBulkMail(self, iListMail, iMailbox, iUser):
        """Bulk save of the mail"""
        for m in iListMail:
            self.SaveMail(m, iMailbox, iUser)
    
    def SaveMail(self, iMail, iMailbox, iUser):
        mailbox = self.session.query(Mailbox).filter_by(alias = iMailbox.alias).\
                  filter(Mailbox.users.has(User.username == iUser.username))
        if mailbox:
            mailbox[0].mail.append(iMail)
            self.session.save(iMail)
        else:
            raise EmptyObjectException('User named %s does not have a mailbox named %s' % (iUser.username,
                                                                                           iMailbox.alias))
            
    def SaveBulkAttachment(self, iListMailAttachment):
        """Bulk save of the attachment"""
        for mail, att, in iListMailAttachment:
            self.SaveAttachment(mail, att)

    def SaveAttachment(self, iAttachment, iMail):
        mail = self.session.query(Mail).filter_by(body = iMail.body, orig_date = iMail.orig_date,
                                                  from_who = iMail.from_who, state = iMail.state)
        if mail:
            mail[0].attachment.append(iAttachment)
            self.session.save(iAttachment)
        else:
            raise EmptyObjectException('Mail does not exit')            

    def SaveAddressBook(self, iAddressBook, iUser):
        user = self.session.query(User).filter_by(username = iUser.username)
        if user:
            user[0].addressbook.append(iAddressBook)
            self.session.save(iAddressBook)
        else:
            raise EmptyObjectException('User named %s does not exist in the database' % iUser.username)
        
    def SaveAddress(self, iAddress, iAddressBook, iUser):
        addressbook = self.session.query(AddressBook).filter_by(alias = iAddressBook.alias).\
                      filter(AddressBook.users.has(User.username == iUser.username))
        if addressbook:
            addressbook[0].address.append(iAddress)
            self.session.save(iAddress)
        else:
            raise EmptyObjectException('User name %s does not have an addressbook named %s' % (iUser.username,
                                                                                               iAddressbok.alias))
    def close(self):
        """Close the session and commit all transactions to the database"""
        self.session.commit()
        self.session.close()

if __name__ == '__main__':
    import sys
    b = 'am'
    if b == 's':
        broker = MailBroker()
        
        # Objects to test database with
        user = User('danielo', 'hippie')
        owner = Owner('danielo', 'hippie')
        control = Control(True, True, True)
        mailbox = Mailbox('home', True)
        mail = Mail(datetime.datetime.now(), 'KuplaCreator@kuplaproject.com', 'Hello', 'Draft')
        domain = Domain('danielo', 'umail.ucsb.edu', 'hippie')
        addressbook = AddressBook('Contacts', True)
        address = Address('KuplaCreator@kuplaproject.com', 'Hippie 3')
        first_objects = [user, owner, control, mailbox, mail, domain, addressbook, address]
        
        # Testing Save
        try:
            broker.SaveOwner(owner)
        except Exception, e:
            print e
        broker.SaveUser(user, owner)
        try:
            broker.SaveControl(control, owner)
        except Exception, e:
            print e
        broker.SaveMailbox(mailbox, user)
        broker.SaveDomain(domain, mailbox, user)
        broker.SaveMail(mail, mailbox, user)
        broker.SaveAddressBook(addressbook, user)
        broker.SaveAddress(address, addressbook, user)
        
        # Testing Has
        if broker.HasOwner(owner) and broker.HasUser(user) and broker.HasControl(control) \
           and broker.HasMailbox(mailbox, user) and broker.HasDomain(domain, mailbox, user) \
           and broker.HasMail(mail, mailbox, user) and broker.HasAddressBook(addressbook, user) \
           and broker.HasAddress(address, addressbook, user):
            print 'Success: Has'
        else:
            print 'Failure: Has'
        
        # Testing Get
        print broker.GetOwner()[0]
        print broker.GetUser()[0]
        print broker.GetSuperUser()[0]
        print broker.GetControl()[0]
        print broker.GetMailbox()[0]
        print broker.GetUserMailbox(user)[0]
        print broker.GetAddressBook()[0]
        print broker.GetUserAddressBook(user)[0]
        print broker.GetMail()[0]
        print broker.GetMailboxMail(mailbox)[0]
        print broker.GetDomain()[0]
        print broker.GetMailboxDomain(mailbox)[0]
        print broker.GetAddress()[0]
        print broker.GetAddressBookAddress(addressbook)[0]
        broker.close()
    else:
        from twisted.internet import reactor
        def TestData(data):
            print data
        def CreateMail(mailbox, **kwargs):
            broker = kwargs['broker']
            print mailbox.id
            mail = [Mail(-1, "", "", "", 0,"", "", "", "", "", "", "", "", "", "", "", "", mailbox.id) for v in range(50)]
            return broker.SaveMailList(mail).addErrback(TestData)
        def CreateMailbox(broker):
            m = Mailbox(0, '', 0, '', '', '', 0, '', '', 0, True, '', 0, True)
            print m
            return broker.SaveMailbox(m).addCallback(lambda _: broker.GetMailbox(m)).addErrback(TestData)
        def TestUpdate(mailbox, **kwargs):
            broker = kwargs['broker']
            print mailbox
            m2 = Mailbox(*mailbox.get())
            m2.size = 255
            return broker.UpdateMailbox(mailbox, m2).addCallback(lambda _: broker.GetMailbox(m2)).addErrback(TestData)
        def TestMailUpdate(mailbox, **kw):
            def Modify(m, **kw):
                old_mail = Mail(*m.get())
                m.local_flags = '\\Deleted'
                mailbox = kw['mailbox']
                broker = kw['broker']
                return broker.UpdateMail(old_mail, m, mailbox)
            broker = kw['broker']
            m =  Mail(-1, "", "", "", 0,"", "", "", "", "", "", "", "", "", "", "", "", mailbox.id) 
            d = broker.SaveMail(m)
            d.addCallback(lambda _: broker.GetMail(m))
            d.addCallback(Modify, broker = broker, mailbox = mailbox)
        broker = AsynchronousMailBroker("sqlite:///mail24.db")
        d = broker.startup()
        d.addCallback(lambda _: broker.GetData()).addErrback(TestData)
        d.addBoth(TestData)
        d.addCallback(lambda _: CreateMailbox(broker))
        d.addCallback(CreateMail, broker = broker)
        d.addBoth(TestData)
        d.addCallback(lambda _: broker.GetData())
        d.addCallback(TestData)
        d.addCallback(lambda _: CreateMailbox(broker))
        d.addCallback(TestUpdate, broker = broker)
        d.addCallback(lambda _: broker.GetData())
        d.addBoth(TestData)
        d.addCallback(lambda _: CreateMailbox(broker))
        d.addCallback(TestMailUpdate, broker = broker)
        d.addCallback(lambda _: broker.GetData())
        d.addBoth(TestData)
        d.addCallback(lambda _: broker.DeleteAll())
        d.addCallback(lambda _: broker.shutdown())
        d.addCallback(lambda _: reactor.stop())
        reactor.run()